package com.tyrone.algorithm.graph;

import com.tyrone.algorithm.linear.Queue;
import com.tyrone.algorithm.priority.IndexMinPriorityQueue;

/**
 * Prim算法计算最小生成树
 */
public class PrimMST {
    private Edge[] edgeTo;                  //索引代表顶点，值表示当前顶点和最小生成树之间的最短边
    private double[] distTo;                //索引代表顶点，值表示当前顶点和最小生成树之间的最短边的权重
    private boolean[] marked;               //索引代表顶点，如果当前顶点已经在树中，则值为true，否则为false
    private IndexMinPriorityQueue<Double> pq;       //存放树中顶点与非树中顶点之间的有效横切边


    public PrimMST(EdgeWeightedGraph G ){
        edgeTo = new Edge[G.V()];
        distTo = new double[G.V()];
        for (int i = 0; i < distTo.length; i++) {
            distTo[i]=Double.POSITIVE_INFINITY;
        }
        marked=new boolean[G.V()];
        pq=new IndexMinPriorityQueue(G.V());
        distTo[0]=0.0;
        pq.insert(0,0.0);
        while (!pq.isEmpty()){
            visit(G,pq.delMin());
        }
    }

    /**
     * 将顶点v添加到最小生成树中，并且更新数据
     * @param G
     * @param v
     */
    private void visit(EdgeWeightedGraph G, int v) {
        marked[v]=true;
        for (Edge e : G.adj(v)) {
            int w = e.other(v);
            if (marked[w]){
                continue;
            }
            if(e.weight()<distTo[w]){
                edgeTo[w]=e;
                distTo[w]=e.weight();
                if (pq.contains(w)){
                    pq.changeItem(w,e.weight());
                }else {
                    pq.insert(w,e.weight());
                }
            }
        }
    }

    /**
     * 获取最小生成树的所有边
     * @return
     */
    public Queue<Edge> edges() {
        Queue<Edge> edgeQueue = new Queue<>();
        for (int i = 0; i < marked.length; i++) {
            if (edgeTo[i]!=null){
                edgeQueue.enqueue(edgeTo[i]);
            }
        }
        return edgeQueue;
    }
}
